import asyncio
from autogen_ext.models.openai import OpenAIChatCompletionClient
from autogen_core.models import ModelFamily
from autogen_agentchat.agents import AssistantAgent
from autogen_agentchat.teams import RoundRobinGroupChat
from autogen_agentchat.conditions import TextMentionTermination
from autogen_agentchat.ui import Console
from utils.get_llm_config import get_llm_config
from bdtime import tt


async def main():
    # 用户需求
    # project_goal = input("请输入您的开发需求: ")
    project_goal = "写一个贪吃蛇小游戏, 用html完成"

    llm_config = get_llm_config()

    model_client = OpenAIChatCompletionClient(
        **llm_config,
        model_info={
            "vision": False,
            "function_calling": False,
            "json_output": False,
            "family": ModelFamily.UNKNOWN,
            "structured_output": True,
        },
        timeout=600,
        max_retries=3,
    )

    def get_doc(doc_path: str):
        """读取指定路径, 转为str"""
        with open(doc_path, "r", encoding="utf-8") as f:
            return f.read()

    example_prd_doc = get_doc("docs/example_prd.txt")

    prd_system_message = """你是产品需求分析师。根据用户需求生成完整的PRD文档.
<example_prd_doc>
{example_prd_doc}
</example_prd_doc>
按照示例, 输出专业的markdown格式PRD文档。""".format(example_prd_doc=example_prd_doc)

    prd_writer = AssistantAgent(
        name="prd_writer",
        model_client=model_client,
        system_message=prd_system_message,
    )

    print(f"--- tt.now: {tt.now()} ----- start real_prd_doc\n")
    prd_result = await prd_writer.run(task=project_goal)
    real_prd_doc = prd_result.messages[-1].content
    print(f"=== tt.now: {tt.now()} ----- end real_prd_doc\n")

    from prompts.prompts import BASE_PLANNING_PROMPT

    planning_system_message = BASE_PLANNING_PROMPT
    planning_system_message = planning_system_message.replace("{project_goal}", project_goal)
    planning_system_message = planning_system_message.replace("{requirement_document}", real_prd_doc)

    task_manager = AssistantAgent(
        name="task_manager",
        model_client=model_client,
        system_message=planning_system_message
    )

    reviewer_system_message = """你是质量审核专家。审核PRD文档和任务拆解的完整性、合理性，提供改进建议。
    <评估标准参考>
    {example_evaluation_rule}
    </评估标准参考>
    若判定任务拆分结果质量合格, 不需要修改了, 则最后输出 'COMPLETE' 表示工作流结束。"""
    example_evaluation_rule = get_doc("docs/规划评估标准参考.txt")
    reviewer_system_message = reviewer_system_message.replace("{example_evaluation_rule}", example_evaluation_rule)
    reviewer = AssistantAgent(
        name="reviewer",
        model_client=model_client,
        system_message=reviewer_system_message
    )

    termination = TextMentionTermination("COMPLETE")

    # members = [prd_writer, task_manager, reviewer]
    members = [task_manager, reviewer]
    team = RoundRobinGroupChat(members, termination_condition=termination)

    # 运行工作流
    print("🚀 开始生成PRD和拆解任务...")

    task_prompt = f"""
    用户开发需求：{project_goal}

    请按顺序完成：
    1. TASK_MANAGER: 拆解为开发子任务  
    2. REVIEWER: 审核并确认完成, 若审核不通过, 则给出修改意见, 并返回上一步, 让 TASK_MANAGER 按照修改意见重新拆分.
    """
    # await Console(team.run_stream(task=task_prompt))

    # Run the team again without a task to continue the previous task.
    stream = team.run_stream(task=task_prompt)
    async for message in stream:
        print(f"\n--- tt.now: {tt.now()}")
        print(message)
        print("=" * 20)


if __name__ == "__main__":
    asyncio.run(main())
